# %BANNER_BEGIN%
# ---------------------------------------------------------------------
# %COPYRIGHT_BEGIN%
#
#  Magic Leap, Inc. ("COMPANY") CONFIDENTIAL
#
#  Unpublished Copyright (c) 2020
#  Magic Leap, Inc., All Rights Reserved.
#
# NOTICE:  All information contained herein is, and remains the property
# of COMPANY. The intellectual and technical concepts contained herein
# are proprietary to COMPANY and may be covered by U.S. and Foreign
# Patents, patents in process, and are protected by trade secret or
# copyright law.  Dissemination of this information or reproduction of
# this material is strictly forbidden unless prior written permission is
# obtained from COMPANY.  Access to the source code contained herein is
# hereby forbidden to anyone except current COMPANY employees, managers
# or contractors who have executed Confidentiality and Non-disclosure
# agreements explicitly covering such access.
#
# The copyright notice above does not evidence any actual or intended
# publication or disclosure  of  this source code, which includes
# information that is confidential and/or proprietary, and is a trade
# secret, of  COMPANY.   ANY REPRODUCTION, MODIFICATION, DISTRIBUTION,
# PUBLIC  PERFORMANCE, OR PUBLIC DISPLAY OF OR THROUGH USE  OF THIS
# SOURCE CODE  WITHOUT THE EXPRESS WRITTEN CONSENT OF COMPANY IS
# STRICTLY PROHIBITED, AND IN VIOLATION OF APPLICABLE LAWS AND
# INTERNATIONAL TREATIES.  THE RECEIPT OR POSSESSION OF  THIS SOURCE
# CODE AND/OR RELATED INFORMATION DOES NOT CONVEY OR IMPLY ANY RIGHTS
# TO REPRODUCE, DISCLOSE OR DISTRIBUTE ITS CONTENTS, OR TO MANUFACTURE,
# USE, OR SELL ANYTHING THAT IT  MAY DESCRIBE, IN WHOLE OR IN PART.
#
# %COPYRIGHT_END%
# ----------------------------------------------------------------------
# %AUTHORS_BEGIN%
#
#  Originating Authors: Paul-Edouard Sarlin
#
# %AUTHORS_END%
# --------------------------------------------------------------------*/
# %BANNER_END%

import torch
import yaml
import cv2
import argparse



from DataLoader import create_dataloader
from Detectors import create_detector
from Matchers import create_matcher

from models.utils import AverageTimer
from utils.tools import *


class Matching(torch.nn.Module):
    """ Image Matching Frontend (SuperPoint + SuperGlue) """
    def __init__(self, config={}):
        super().__init__()
        # create detector
        self.detector = create_detector(config["detector"])
        # create matcher
        self.matcher = create_matcher(config["matcher"])


        # self.superpoint = SuperPoint(config.get('superpoint', {}))
        # self.superglue = SuperGlue(config.get('superglue', {}))

    def forward(self, data):
        """ Run SuperPoint (optionally) and SuperGlue
        SuperPoint is skipped if ['keypoints0', 'keypoints1'] exist in input
        Args:
          data: dictionary with minimal keys: ['image0', 'image1']
        """
        pred = {'ref': None, 'cur': None}

        # TODO 这块的显存存占用有待优化
        # Extract SuperPoint (keypoints, scores, descriptors) if not provided
        if 'keypoints0' not in data:
            # pred0 = self.superpoint({'image': data['image0']})
            # pred = {**pred, **{k+'0': v for k, v in pred0.items()}}
            pred['ref'] = self.detector(data['image0'])
            # pred0 = self.detector(data['image0'])
            # pred = {**pred, **{k+'0': v for k, v in pred0.items()}}
        if 'keypoints1' not in data:
            # pred1 = self.superpoint({'image': data['image1']})
            # pred = {**pred, **{k+'1': v for k, v in pred1.items()}}
            pred['cur'] = self.detector(data['image1'])
            # pred0 = self.detector(data['image1'])
            # pred = {**pred, **{k+'1': v for k, v in pred0.items()}}



        # Batch all features
        # We should either have i) one image per batch, or
        # ii) the same number of local features for all images in the batch.
        # data = {**data, **pred}
        #
        # for k in data:
        #     if isinstance(data[k], (list, tuple)):      # 将 list, tuple类转换为tensor  TODO 这部分或许可以优化
        #         data[k] = torch.stack(data[k])

        # Perform the matching
        # pred = {**pred, **self.superglue(data)}     # TODO pred这部分只用了matches0和matching_score0，matches1和matching_score1或可删去

        matches = self.matcher(pred)

        return matches


def get_args():
    parser = argparse.ArgumentParser(
        description='Image pair matching and pose evaluation with SuperGlue',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--idx', type=int, default=0,
        help='idx of config files')

    # parser.add_argument('--step', type=int, default=1,
    #     help='step between frames')

    return parser.parse_args()

if __name__ == "__main__":
    # 这一句至关重要！！！ 能节省至少一半显存！！！
    torch.set_grad_enabled(False)

    opt = get_args()
    print(opt)

    yamls = ["params/superpoint_supergluematch.yaml", "params/superpoint_flannmatch.yaml",
             "params/sift_flannmatch.yaml", "params/orb_brutematch.yaml"]
    # config = "params/kitti_superpoint_supergluematch.yaml"
    config = yamls[opt.idx]

    with open(config, 'r') as f:
        config = yaml.safe_load(f)
    # create dataloader
    step = config['dataset']['step']
    loader = create_dataloader(config["dataset"])
    pair_num = len(loader)

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    matching = Matching(config).eval().to(device)

    imgs = {}
    timer = AverageTimer(newline=True)

    print ("OK")
    for i, img in enumerate(loader):
        # imgs["cur"] = torch.from_numpy(img/255.).float()[None, None].to(matcher.device)
        imgs["cur"] = img
        # kptdescs["cur"] = detector(img)
        timer.update('load_img')

        if i >= 1:
            matches = matching({'image0': imgs['ref'], 'image1': imgs['cur']})
            timer.update('matcher')

            if opt.idx < 2:
                img = plot_matches(imgs['ref'], imgs['cur'],
                               matches['ref_keypoints'].cpu().detach().numpy()[0::5], matches['cur_keypoints'].cpu().detach().numpy()[0::5],
                               matches['match_score'].cpu().detach().numpy()[0::5], layout='lr')
            else:
                img = plot_matches(imgs['ref'], imgs['cur'],
                                   matches['ref_keypoints'][0::5], matches['cur_keypoints'][0::5],
                                   matches['match_score'][0::5], layout='lr')
            cv2.imshow("track", img)
            if cv2.waitKey(1) == 27:     # 按 ‘ESC'结束运行
                break
            timer.update('viz')
            timer.print('Keypoint num {} \t Finished pair {:5} of {:5}'.format(len(matches['ref_keypoints']), i*step, pair_num))

            # 用不着了（其实也不起啥作用）
            # matches.clear()
            # torch.cuda.empty_cache()


        # kptdescs["ref"], imgs["ref"] = kptdescs["cur"], imgs["cur"]
        # kptdescs["ref"] = kptdescs["cur"]
        imgs["ref"] = imgs["cur"]


